package hn.cch.example.lts17;

import com.google.gson.Gson;
import com.google.gson.JsonObject;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.HasCollectionReq;
import io.milvus.v2.service.collection.request.LoadCollectionReq;
import io.milvus.v2.service.collection.request.ReleaseCollectionReq;
import io.milvus.v2.service.database.request.CreateDatabaseReq;
import io.milvus.v2.service.index.request.CreateIndexReq;
import io.milvus.v2.service.partition.request.ListPartitionsReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.BaseVector;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.SearchResp;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class MilvusDemo {

    public static void main(String[] args) throws InterruptedException {

        String uri = "https://in03-4a005256ce9931b.serverless.aws-eu-central-1.cloud.zilliz.com";
        String token = "a09138af5b8ab9e8d392ecdf6d58ddde6dc5d6b28f03c842c786a8b1430a12713582eebb6287878fa3c7ebf44b3282e0c7c3d672";
        ConnectConfig config = ConnectConfig.builder().uri(uri).token(token).build();
        MilvusClientV2 client = new MilvusClientV2(config);
        String databaseName = "hangcchhn";
        client.createDatabase(CreateDatabaseReq.builder()
            .databaseName(databaseName)
            .build());
        String collectionName = "image_vector_db";
        client.useDatabase(databaseName);
        boolean exists = client.hasCollection(HasCollectionReq.builder()
            .collectionName(collectionName)
            .build());


        List<String> listPartitions = client.listPartitions(ListPartitionsReq.builder()
            .collectionName(collectionName)
            .build());


        CreateCollectionReq.CollectionSchema schema = MilvusClientV2.CreateSchema();

        String fieldName = "float_vector";
        schema.addField(AddFieldReq.builder()
            .fieldName(fieldName)
            .dataType(DataType.FloatVector)
            .dimension(5)
            .build());

        IndexParam indexParam = IndexParam.builder()
            .fieldName(fieldName)
            .indexType(IndexParam.IndexType.IVF_FLAT)
            .metricType(IndexParam.MetricType.L2)
            .extraParams(Map.of("n", 128))
            .build();
        client.createIndex(CreateIndexReq.builder()
            .collectionName(collectionName)
            .indexParams(List.of(indexParam))
            .build());

        Gson gson = new Gson();
        List<Map<String, Object>> rows = new ArrayList<>();
        List<JsonObject> jsonData = rows.stream()
            .map(m -> gson.toJsonTree(m).getAsJsonObject())
            .toList();

        client.insert(InsertReq.builder()
            .collectionName(collectionName)
            .data(jsonData)
            .build());

        client.loadCollection(LoadCollectionReq.builder()
            .collectionName(collectionName)
            .build());

        float[] array = new float[]{0.15f, 0.25f, 0.35f, 0.45f, 0.55f};
        List<BaseVector> vectorData = List.of(new FloatVec(array));
        SearchReq searchReq = SearchReq.builder()
            .collectionName(collectionName)
            .annsField(fieldName)
            .metricType(IndexParam.MetricType.L2)
            .data(vectorData)
            .limit(3)
            .searchParams(Map.of("n", 10))
            .outputFields(List.of("label"))
            .build();
        SearchResp searchResp = client.search(searchReq);

        client.createCollection(CreateCollectionReq.builder()
            .collectionName(collectionName)
            .collectionSchema(schema)
            .build());

        client.releaseCollection(ReleaseCollectionReq.builder()
            .collectionName(collectionName)
            .build());
        client.close();
    }

}
